package grid_operation;

import com.google.gson.Gson;
import gdal_grid.Raster2Vector;
import org.gdal.gdal.*;
import org.gdal.gdalconst.gdalconstConstants;
import utils.SelectFile;

import java.io.File;
import java.util.*;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

public class GridOperation {
    public static void main(String[] a) {
        String tifLinDI = "E:\\Test\\database1\\result\\林地地类.tif";
        String tifXiongjin = "E:\\Test\\database1\\result\\平均胸径.tif";
        String tifPu = "E:\\Test\\database1\\result\\坡度.tif";
        String tifPuWei = "E:\\Test\\database1\\result\\坡位.tif";
        String tifPuXiang = "E:\\Test\\database1\\result\\坡向.tif";
        String tifTucen = "E:\\Test\\database1\\result\\土层厚度.tif";
        String tifTuCenMax = "E:\\TestInput\\TestInput2\\tucenghoudu.tif";
        String tifLanUse = "E:\\TestInput\\TestInput2\\landuse.tif";
        String tifDILEID = "E:\\Test\\WaterResources_Clip.tif";
        String tifTest = "C:\\Users\\PS2018\\Documents\\WXWork\\1688850924749243\\Cache\\File\\2021-03\\1616122961611(1).tif";

        //       GridBean gridBeanDi = new GridBean(new File(tifLinDI), OperationSymbol.EQUAL, 111);
//        GridBean gridBeanPu = new GridBean(new File(tifPu), OperationSymbol.LESS_THAN, 50);
//        GridBean gridBeanPuWei = new GridBean(new File(tifPuWei), OperationSymbol.EQUAL, 4);
        //GridBean gridBeanTucen = new GridBean(new File(tifTuCenMax), OperationSymbol.GREATER_THAN, 20);
        GridBean gridBeanLanUse = new GridBean(new File(tifTest), OperationSymbol.EQUAL, 111);

        List<GridBean> gridBeanList = new ArrayList<>();
        //    gridBeanList.add(gridBeanDi);
//        gridBeanList.add(gridBeanPu);
        //    gridBeanList.add(gridBeanPuWei);
        //gridBeanList.add(gridBeanTucen);
        gridBeanList.add(gridBeanLanUse);
        File out = new File("E:\\Test\\result\\result" + System.currentTimeMillis() + ".tif");
        gridOper(gridBeanList, out);
    }

    //核心线程数
    static int corePoolSize = Runtime.getRuntime().availableProcessors();
    // 池中允许的最大线程数
    static final int maxPoolSize = corePoolSize * 2 + 1;
    // 线程池中的空闲线程所能持续的最长时间
    static final long keepAliveTime = 0L;
    // 时间单位
    static final TimeUnit unit = TimeUnit.MILLISECONDS;
    /**
     * 等待队列大小
     */
    private static int QUEUE_SIZE = 1024;

    /**
     * 栅格运算
     * @param gridBeans 栅格计算对象
     * @param outFile 输出路径
     */
    public static void gridOper(List<GridBean> gridBeans, File outFile) {
        long time1 = System.currentTimeMillis();

        ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(corePoolSize,
                maxPoolSize, keepAliveTime, unit, new ArrayBlockingQueue<Runnable>(QUEUE_SIZE),
                new ThreadPoolExecutor.CallerRunsPolicy());

        if (gridBeans == null || gridBeans.isEmpty()) {
            return;
        }

        gdal.AllRegister();
        gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES");
        gdal.SetConfigOption("SHAPE_ENCODING", "UTF-8");

        //打开tif文件
        List<Dataset> tifDataSet = new ArrayList<>();
        Map<Dataset, GridBean> datasetMap = new HashMap<>();
        for (GridBean gridBean : gridBeans) {
            Dataset hDataset = gdal.Open(gridBean.getGridFile().getPath(), gdalconstConstants.GA_ReadOnly);
            tifDataSet.add(hDataset);
            datasetMap.put(hDataset, gridBean);
        }

        //计算最小栅格边界
        int iXSize = Integer.MAX_VALUE;
        int iYSize = Integer.MAX_VALUE;
        Dataset xMinDataset = tifDataSet.get(0);
        Dataset yMinDataset = tifDataSet.get(0);
        for (Dataset dataset : tifDataSet) {
            int x = dataset.getRasterXSize();
            int y = dataset.getRasterYSize();

            if (x <= iXSize) {
                iXSize = x;
                xMinDataset = dataset;
            }

            if (y <= iYSize) {
                iYSize = y;
                yMinDataset = dataset;
            }
        }

        List<Band> tifBand = new ArrayList<>();
        Map<Band, GridBean> bandMap = new HashMap<>();
        for (Dataset dataset : tifDataSet) {
            Band band = dataset.GetRasterBand(1);
            tifBand.add(band);
            bandMap.put(band, datasetMap.get(dataset));
        }
        //创建结果图像
        Driver driver = gdal.GetDriverByName("GTIFF");
        //六参数信息
        double[] geoTransform = xMinDataset.GetGeoTransform();
        geoTransform[3] = yMinDataset.GetGeoTransform()[3];

        File outDir = new File(outFile.getParent());
        if (!outDir.exists()) {
            outDir.mkdirs();
        }

        //生成输出栅格文件
        File tempFile = new File(outDir.getPath() + File.separator + System.currentTimeMillis() + "tif");
        Dataset outputDs = driver.Create(tempFile.getPath(), iXSize, iYSize, 1, gdalconstConstants.GDT_Byte);
        outputDs.SetGeoTransform(geoTransform);
        outputDs.SetProjection(xMinDataset.GetProjection());

        Band bandResult = outputDs.GetRasterBand(1);
        bandResult.SetNoDataValue(255);

        for (int i = 0; i < iYSize; i++) {
            threadPoolExecutor.execute(new OperationTaskRunnable(iXSize, tifBand, bandMap, bandResult, i));
        }
        threadPoolExecutor.shutdown();
        while (true) {//等待所有线程执行完毕
            if (threadPoolExecutor.isTerminated()) {
                System.out.println("所有的子线程都结束了！");
                break;
            }
        }

        long time2 = System.currentTimeMillis();
        System.out.println("计算执行时间：" + (time2 - time1));

        //压缩文件
        List<String> op = new ArrayList<>();
        op.add("TILED=YES");
        op.add("COMPRESS=LZW");
        Dataset copyDataset = driver.CreateCopy(outFile.getPath(), outputDs, 1, new Vector(op));
        outputDs.delete();
        tempFile.deleteOnExit();

        long time3 = System.currentTimeMillis();
        System.out.println("压缩执行时间：" + (time3 - time2));

        //分割文件
        List<Dataset> datasetList = cuttingTif(copyDataset, outFile, driver);
        copyDataset.delete();
        long time4 = System.currentTimeMillis();
        System.out.println("裁切执行时间：" + (time4 - time3));


        org.gdal.ogr.Driver driverData = SelectFile.driverInit();
        Raster2Vector.dataSource = driverData.CreateDataSource(outDir + "//" + outFile.getName() + "Un.shp");

        ThreadPoolExecutor threadPoolExecutor2 = new ThreadPoolExecutor(corePoolSize,
                maxPoolSize, keepAliveTime, unit, new ArrayBlockingQueue<Runnable>(QUEUE_SIZE),
                new ThreadPoolExecutor.CallerRunsPolicy());
        for (int i = 0; i < datasetList.size(); i++) {
            int finalI = i;
            threadPoolExecutor2.execute(() -> {
                Raster2Vector.raster2V(datasetList.get(finalI), outDir + "//" + outFile.getName() + finalI + ".shp", driverData);
            });
        }
        threadPoolExecutor2.shutdown();
        while (true) {//等待所有线程执行完毕
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            if (threadPoolExecutor2.isTerminated()) {
                System.out.println("所有的子线程都结束了！");
                break;
            }
        }

        //删除临时文件
        for (int i = 0; i < cutIndex * cutIndex; i++) {
            String outputPath = outFile.getParent() + File.separator + "test" + i + ".tif";
            File file = new File(outputPath);
            file.delete();

            for (int k = 0; k < suffixs.length; k++) {
                String outPutShapePath = outDir + "//" + outFile.getName() + i + suffixs[k];
                File shapeFile = new File(outPutShapePath);
                shapeFile.delete();
            }
        }

        System.out.println("总执行时间：" + (System.currentTimeMillis() - time1));
    }

    private static String[] suffixs = new String[]{
            ".shp",
            ".dbf",
            ".prj",
            ".shx",
    };
    private static int cutIndex = 30;

    private static List<Dataset> cuttingTif(Dataset copyDataset, File outFile, Driver driver) {

        int ix = copyDataset.getRasterXSize() / cutIndex;
        int iy = copyDataset.getRasterYSize() / cutIndex;

        List<Dataset> cutDataSet = new ArrayList<>();
        int index = 0;
        for (int i = 0; i < cutIndex; i++) {
            for (int h = 0; h < cutIndex; h++) {
                double[] geoTransform = copyDataset.GetGeoTransform();
                String outputPath = outFile.getParent() + File.separator + "test" + index + ".tif";
                Dataset outputDs = driver.Create(outputPath, ix, iy, 1, gdalconstConstants.GDT_Byte);
                geoTransform[0] = geoTransform[0] + ix * geoTransform[1] * h + iy * geoTransform[2] * h;
                geoTransform[3] = geoTransform[3] + ix * geoTransform[4] * i + iy * geoTransform[5] * i;

                outputDs.SetGeoTransform(geoTransform);
                outputDs.SetProjection(copyDataset.GetProjection());

                Band band = copyDataset.GetRasterBand(1);
                Band bandResult = outputDs.GetRasterBand(1);
                bandResult.SetNoDataValue(255);
                for (int k = (i * iy); k < ((i + 1) * iy); k++) {
                    int[] cache = new int[ix];
                    band.ReadRaster(ix * h, k, ix, 1, cache);

                    synchronized (GridOptions.class) {
                        bandResult.WriteRaster(0, (k - (i * iy)), ix, 1, cache);
                        bandResult.FlushCache();
                    }
                }
                cutDataSet.add(outputDs);
                index++;
            }
        }

        return cutDataSet;
    }

    static class OperationTaskRunnable implements Runnable {
        private int iXSize;
        private List<Band> tifBand;
        private Map<Band, GridBean> bandMap;
        private Band bandResult;
        private int i;
        private Gson gson = new Gson();

        public OperationTaskRunnable(int iXSize, List<Band> tifBand, Map<Band, GridBean> bandMap, Band bandResult, int i) {
            this.iXSize = iXSize;
            this.tifBand = tifBand;
            this.bandMap = bandMap;
            this.bandResult = bandResult;
            this.i = i;
        }

        @Override
        public void run() {
            List<int[]> cacheList = new ArrayList<>();
            Map<int[], GridBean> cacheMap = new HashMap<>();
            synchronized (OperationTaskRunnable.class) {
                for (Band band : tifBand) {
                    Double[] noData = new Double[1];
                    band.GetNoDataValue(noData);
                    if (noData[0] == null) {
                        noData[0] = Double.valueOf(255);
                    }
                    int[] cache = new int[iXSize];
                    band.ReadRaster(0, i, iXSize, 1, cache);
                    System.out.println(gson.toJson(cache));
                    cacheList.add(cache);
                    GridBean gridBean = bandMap.get(band);
                    gridBean.setNoData(noData[0]);
                    cacheMap.put(cache, gridBean);
                }
            }

            int[] cacheWr = new int[iXSize];
            for (int k = 0; k < iXSize; k++) {
                List<Integer> cacheXList = new ArrayList<>();
                Map<Integer, GridBean> intMap = new HashMap<>();
                for (int[] is : cacheList) {
                    cacheXList.add(is[k]);
                    intMap.put(is[k], cacheMap.get(is));
                }
                boolean isSuccess = operation(cacheXList, intMap);

                if (isSuccess) {
                    cacheWr[k] = 1;
                } else {
                    cacheWr[k] = 255;
                }
            }

            synchronized (GridOptions.class) {
                bandResult.WriteRaster(0, i, iXSize, 1, cacheWr);
                bandResult.FlushCache();
            }
        }
    }

    private static boolean operation(List<Integer> cacheList, Map<Integer, GridBean> intMap) {
        List<Boolean> result = new ArrayList<>();
        for (int cache : cacheList) {
            GridBean gridBean = intMap.get(cache);
            int value = gridBean.relationValue;
            //System.out.println("c = " + cache);
            OperationSymbol operationSymbol = gridBean.operationSymbol;
            Double noData = gridBean.noData;

            boolean is = false;
            if (cache != noData.intValue()) {
                switch (operationSymbol) {
                    case EQUAL:
                        if (value == cache) {
                            is = true;
                        }
                        break;
                    case GREATER_THAN:
                        if (cache > value) {
                            is = true;
                        }
                        break;
                    case LESS_THAN:
                        if (cache < value) {
                            is = true;
                        }
                        break;
                    case LESS_EQUAL_THAN:
                        if (cache <= value) {
                            is = true;
                        }
                        break;
                    case GREATER_EQUAL_THAN:
                        if (cache >= value) {
                            is = true;
                        }
                        break;
                }
            }
            result.add(is);
        }

        return !result.contains(false);
    }

    static class GridBean {
        private File gridFile; //tif文件
        private OperationSymbol operationSymbol; //运算关系
        private int relationValue; //运算值
        private double noData;

        public GridBean(File gridFile, OperationSymbol operationSymbol, int relationValue) {
            this.gridFile = gridFile;
            this.operationSymbol = operationSymbol;
            this.relationValue = relationValue;
        }

        public double getNoData() {
            return noData;
        }

        public void setNoData(double noData) {
            this.noData = noData;
        }

        public File getGridFile() {
            return gridFile;
        }

        public void setGridFile(File gridFile) {
            this.gridFile = gridFile;
        }

        public OperationSymbol getOperationSymbol() {
            return operationSymbol;
        }

        public void setOperationSymbol(OperationSymbol operationSymbol) {
            this.operationSymbol = operationSymbol;
        }

        public int getRelationValue() {
            return relationValue;
        }

        public void setRelationValue(int relationValue) {
            this.relationValue = relationValue;
        }
    }

    enum OperationSymbol {
        GREATER_THAN, //大于
        GREATER_EQUAL_THAN, //大于等于
        LESS_THAN, //小于
        LESS_EQUAL_THAN, //小于等于
        EQUAL //等于
    }
}
